{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬]                                     Elapsed time: 7s (7424.4 it/s)            \n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Gray-Scott reaction diffusion simulation in 2D\n",
    "\"\"\"\n",
    "# Credits:\n",
    "#https://k3d-jupyter.readthedocs.io/en/latest/showcase/gray_scott.html\n",
    "import numpy as np\n",
    "import pycuda.gpuarray as gpuarray\n",
    "from pycuda.compiler import SourceModule\n",
    "import pycuda.driver as cuda\n",
    "\n",
    "#Du, Dv, F, k = 0.16, 0.08, 0.035, 0.065 # Bacteria 1\n",
    "#Du, Dv, F, k = 0.14, 0.06, 0.035, 0.065 # Bacteria 2\n",
    "Du, Dv, F, k = 0.16, 0.08, 0.060, 0.062 # Coral\n",
    "#Du, Dv, F, k = 0.19, 0.05, 0.060, 0.062 # Fingerprint\n",
    "#Du, Dv, F, k = 0.12, 0.08, 0.020, 0.050 # Spirals Dense\n",
    "#Du, Dv, F, k = 0.10, 0.16, 0.020, 0.050 # Spirals Fast\n",
    "#Du, Dv, F, k = 0.16, 0.08, 0.020, 0.055 # Unstable\n",
    "#Du, Dv, F, k = 0.16, 0.08, 0.050, 0.065 # Worms 1\n",
    "#Du, Dv, F, k = 0.16, 0.08, 0.054, 0.063 # Worms 2\n",
    "#Du, Dv, F, k = 0.16, 0.08, 0.035, 0.060 # Zebrafish\n",
    "\n",
    "dt = 0.2\n",
    "size = (512, 512)\n",
    "\n",
    "cuda.init()\n",
    "device = cuda.Device(0)\n",
    "ctx = device.make_context()\n",
    "block_size = 128\n",
    "nx = (size[0]//block_size) * block_size\n",
    "ny = size[1]\n",
    "blocks = nx * ny // block_size\n",
    "u = np.ones( (ny,nx), dtype=np.float32)\n",
    "v = np.zeros((ny,nx), dtype=np.float32)\n",
    "\n",
    "pars = {'nx':nx, 'ny':ny, 'Du':Du, 'Dv':Dv, 'dt':dt, 'F':F, 'k':k}\n",
    "src = \"\"\"\n",
    "    __device__ inline float laplace2d(int idx, float *a) {{\n",
    "      return(a[idx-1] + a[idx+1] + a[idx-{nx}] + a[idx+{nx}] - 4.0f * a[idx] );\n",
    "    }}\n",
    " \n",
    "    __global__ void iterate_RDS(float *a,float *da,float *b,float *db)\n",
    "    {{\n",
    "      int idx = blockDim.x*blockIdx.x + threadIdx.x;\n",
    "      float k = {k}f;  \n",
    "      float F = {F}f;\n",
    "      \n",
    "      if(idx<{nx} || idx>{nx}*({ny}-1)-1) {{\n",
    "          return;\n",
    "      }}\n",
    "      \n",
    "      int x = idx % {nx};\n",
    "      \n",
    "      if(x==0 || x=={nx}-1) {{\n",
    "          return;\n",
    "      }}\n",
    "      \n",
    "      float u = a[idx]; \n",
    "      float v = b[idx];       \n",
    "      da[idx] = u + {dt}f*(-u*v*v + F*(1.0f-u) + {Du}*laplace2d(idx, a));\n",
    "      db[idx] = v + {dt}f*( u*v*v -(F+k)*v     + {Dv}*laplace2d(idx, b));\n",
    "    }}\n",
    "    \"\"\".format(**pars)\n",
    "\n",
    "mod = SourceModule(src)\n",
    "RDSv = mod.get_function(\"iterate_RDS\")\n",
    "\n",
    "r = 100 # initial radius where noise is added\n",
    "u[ny//2-r:ny//2+r,nx//2-r:nx//2+r] = 0.50\n",
    "v[ny//2-r:ny//2+r,nx//2-r:nx//2+r] = 0.25\n",
    "\n",
    "u += 0.05 * np.random.random((ny,nx))\n",
    "v += 0.05 * np.random.random((ny,nx))\n",
    "\n",
    "u_g  = gpuarray.to_gpu(u)\n",
    "du_g = gpuarray.empty_like(u_g)\n",
    "v_g  = gpuarray.to_gpu(v)\n",
    "dv_g = gpuarray.empty_like(v_g)\n",
    "\n",
    "###############################################################\n",
    "from vtkplotter import *\n",
    "embedWindow(False)\n",
    "\n",
    "grid = Grid(resx=size[0]-1, resy=size[1]-1).lw(0).wireframe(False)\n",
    "show(grid, Text(__doc__, c='k'), axes=0, bg='w', interactive=False)\n",
    "\n",
    "pb = ProgressBar(0,50000)\n",
    "for i in pb.range():    \n",
    "    RDSv(u_g,du_g,v_g,dv_g, block=(block_size,1,1), grid=(blocks,1))    \n",
    "    RDSv(du_g,u_g,dv_g,v_g, block=(block_size,1,1), grid=(blocks,1))  \n",
    "    \n",
    "    if i%500 == 0:\n",
    "        v = v_g.get().ravel()\n",
    "        grid.pointColors(v, cmap=\"viridis_r\").show()\n",
    "    pb.print()\n",
    "\n",
    "interactive()\n",
    "closePlotter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
